Made use of the surface_integrals functionality I added while coding the new rotational transform function. The objective will still function in the same way as before if a single rho surface is given as input. However, computing magnetic well over multiple flux surfaces at once offers at least two benefits. It
compute.utils are used then the developer/user doesn't have to worry about it. However, for multiple rho surface grids, the fix isn't as simple. Until that fix is implemented, if the user wants magnetic well computed on a grid with multiple values of rho, for best accuracy they need to set NFP=1 in the grid provided to magnetic well.It seems to have more responsive curves, and it has the same abrupt change as STELLOPT for HELIOTRON. It yields a perfect match for Heliotron.
The curves produced by DESC and STELLOPT for the magnetic well parameter as a function of $\rho$ look different. I think STELLOPT does not implement the well parameter we thought it does. We should just pay attention to the crossing.
The magnetic well parameter computed by DESC matches the sign of the magnetic well parameter computed by STELLOPT. This is technically all we need.
The volume computation is now 100% correct. The solution to the bug in the grid.py class works, which was necessary to correct volume and dv/d$\rho$. An explanation to the solution is given here. Might be useful in case we have run into this bug somewhere else in the code.
All the plots of DESC's intermediate quantities are consistent with each other. For every point below, I have confirmed their validity with code tests and math where applicable. We should be sure to preserve these signs if any modifications are made:
data["V"] on a grid with a sequence of $\rho$ values.B_r and p_r.jnp.mean(f) / jnp.mean($\sqrt{g}$ ) always = jnp.sum(dtdz * f) / dv/d$\rho$ . Helps confirm dv/d$\rho$ is correct among other things.This notebook should make it easy to inspect intermediate variables in the MagneticWell class as well as compute quantities for all the stellarators in a concise manner. Add an entry to the dictionary and everything is automatic. Note this notebook should run pretty quickly. If the plots for many stellarators at once are requested, increase the memory given to jupyter notebook and/or compute on 3 stellarators at a time instead of 5+, and/or reduce plot dpi.
Given this, if the equilibrium solutions DESC and STELLOPT are computing magnetic well parameters on are indeed the same, I think any issue must be in the automatic differentiation of B_r and p_r.
import sys
import os
sys.path.insert(0, os.path.abspath("."))
sys.path.append(os.path.abspath("../"))
%matplotlib inline
import matplotlib.pyplot as plt
import numpy as np
from netCDF4 import Dataset
import pickle
import scipy.io as sio
import desc.io
from desc.compute import compute_geometry, data_index
from desc.equilibrium import Equilibrium
from desc.grid import LinearGrid
from desc.objectives import MagneticWell
from desc.plotting import plot_surfaces
from desc.transform import Transform
DESC version 0.5.1+100.gc4b3314.dirty, using JAX backend, jax version=0.2.25, jaxlib version=0.1.76, dtype=float64 Using device: CPU, with 5.36 GB available memory
np.set_printoptions(precision=4, floatmode="fixed")
class MagneticWellVisual:
"""
To print and plot more (less) quantities, add (remove) them to the dict
returned from MagneticWell.compute().
Everything else is automatic.
"""
def __init__(
self,
name,
eq=None,
use_pickle=True,
rho=np.linspace(1 / 64, 1, 64),
):
"""
Make a MagneticWellVisual from either the provided equilibrium
or the final equilibrium loaded from the output.h5 solution.
Parameters
----------
name : str
Name of the equilibrium.
eq : Equilibrium
The equilibrium.
use_pickle : bool
True to use compute values from a previous run stored in pickle file; False to recompute.
rho : ndarray
The flux surface values to compute the magnetic well on.
"""
self.name = name
if eq is None:
eq = desc.io.load(load_from="../examples/DESC/" + name + "_output.h5")[-1]
self.eq = eq
if use_pickle:
with open(name + " DESC magwell.pkl", "rb") as file:
self.st = pickle.load(file)
self.rho = self.st["volume"][0]
self.has_values = True
else:
# values are tuples of (x[i], y[i]) of plotable data
# y points are cached when compute_plot() is called
self.st = dict()
try:
mat = sio.loadmat(name + "_magwell.mat")
self.st["0. STELLOPT Magnetic Well"] = mat["rho"], mat["magwell"]
rho = mat["rho"]
except FileNotFoundError:
pass
try:
f = Dataset("../iota/edu-vmec/input-iota/wout_" + name + ".nc")
# print(f.variables.keys())
vmec_rho = np.sqrt(
f.variables["phi"] / np.array(f.variables["phi"])[-1]
)
vmec_well = np.asarray(f.variables["DWell"])
self.st["0. VMEC Magnetic Well"] = vmec_rho, vmec_well
except FileNotFoundError:
pass
self.rho = rho
self.has_values = False
def print_values(self, grid=None):
"""
Parameters
----------
grid : LinearGrid
The grid used for a MagneticWell objective.
The default grid=None uses the rho=1 flux surface.
"""
print(self.name)
print(self.eq)
mw = MagneticWell(eq=self.eq, grid=grid)
print("grid.spacing(dr,dt,dz)", mw.grid.spacing[0])
m = mw.compute(
self.eq.R_lmn,
self.eq.Z_lmn,
self.eq.L_lmn,
self.eq.p_l,
self.eq.i_l,
self.eq.c_l,
self.eq.Psi,
)
for key, val in sorted(m.items()):
print(key, val)
self._print_data_V()
print()
def _print_data_V(self):
"""
Print the volume of the stellarator device as computed by data["V"].
Should match V surface integral when the default grid with rho = 1
is used to construct the MagneticWell() object.
"""
grid = LinearGrid(
L=self.eq.L_grid,
M=self.eq.M_grid,
N=self.eq.N_grid,
NFP=self.eq.NFP,
sym=self.eq.sym,
)
R_transform = Transform(
grid, self.eq.R_basis, derivs=data_index["sqrt(g)"]["R_derivs"], build=True
)
Z_transform = Transform(
grid, self.eq.Z_basis, derivs=data_index["sqrt(g)"]["R_derivs"], build=True
)
data = compute_geometry(self.eq.R_lmn, self.eq.Z_lmn, R_transform, Z_transform)
print('data["V"]', data["V"])
# same as other compute_plot but builds plot through computing on single rho surfaces
# def compute_plot(self):
# """Compute and cache MagneticWell.compute() values."""
# if self.has_values:
# return
#
# for i in range(len(self.rho)):
# m = MagneticWell(
# eq=self.eq,
# grid=LinearGrid(
# M=self.eq.M_grid,
# N=self.eq.N_grid,
# NFP=self.eq.NFP,
# sym=self.eq.sym,
# rho=self.rho[i],
# ),
# ).compute(
# self.eq.R_lmn,
# self.eq.Z_lmn,
# self.eq.L_lmn,
# self.eq.p_l,
# self.eq.i_l,
# self.eq.c_l,
# self.eq.Psi,
# )
# for key, val in m.items():
# self.st.setdefault(key, (self.rho, np.empty_like(self.rho)))[1][i] = val
# self.has_values = True
def compute_plot(self):
"""Compute and cache MagneticWell.compute() values."""
if self.has_values:
return
m = MagneticWell(
eq=self.eq,
grid=LinearGrid(
M=self.eq.M_grid,
N=self.eq.N_grid,
NFP=1, # required for accuracy when grid.num_rho != 1
sym=self.eq.sym,
rho=self.rho,
),
).compute(
self.eq.R_lmn,
self.eq.Z_lmn,
self.eq.L_lmn,
self.eq.p_l,
self.eq.i_l,
self.eq.c_l,
self.eq.Psi,
)
for key, val in m.items():
self.st[key] = self.rho, val
self.has_values = True
def save(self):
"""Save computed values to pickle file."""
with open(self.name + " DESC magwell.pkl", "wb") as file:
pickle.dump(self.st, file)
@staticmethod
def _color(key):
"""Return the correct plot color."""
if "0. STELLOPT" in key:
return "tab:green"
if "0. VMEC" in key:
return "tab:brown"
if "1." in key:
return "tab:orange"
if "2." in key:
return "tab:purple"
return "tab:blue"
def plot(self, dpi=300):
"""Plot all quantities from compute_plot() in st."""
fig, ax = plt.subplots(
nrows=3,
ncols=(len(self.st) + 2) // 3,
figsize=(len(self.st) * 3, 15),
dpi=dpi,
)
ax = ax.flatten()
for i, e in enumerate(sorted(self.st.items())):
key, val = e
x, y = val
# edu vmec usually blows up near rho=1. messes up plot scaling if included
if key == "0. VMEC Magnetic Well":
x = x[:-3]
y = y[:-3]
color = MagneticWellVisual._color(key)
ax[i].scatter(x, y, color=color, s=(5 if len(x) > 128 else 10))
ax[i].plot(x, y, color=color)
ax[i].set(xlabel=r"$\rho$", ylabel=key, title=self.name, facecolor="white")
if "Magnetic Well" in key:
if np.any(~np.isclose(y, 0)):
ax[i].set(yscale="symlog" if np.any(y < 0) else "log")
ax[i].axhline(color="tab:red")
ax[i].grid()
def plot_magnetic_wells(self, dpi=300):
"""Plots the magnetic wells together on the same scale."""
fig, ax = plt.subplots(dpi=dpi)
# computed on other codes
if (key := "0. STELLOPT Magnetic Well") in self.st:
color = MagneticWellVisual._color(key)
rho, well = self.st[key]
ax.plot(rho, well, color=color, label="STELLOPT")
ax.scatter(rho, well, color=color, s=1)
if (key := "0. VMEC Magnetic Well") in self.st:
color = MagneticWellVisual._color(key)
rho, well = self.st[key]
# edu vmec usually blows up near rho=1. messes up plot scaling if included
rho, well = rho[:-3], well[:-3]
ax.plot(rho, well, color=color, label="edu-vmec")
ax.scatter(rho, well, color=color, s=1)
# computed with DESC
well = self.st["1. DESC Magnetic Well: M. Landreman eq. 4.19"][1]
color = MagneticWellVisual._color("1.")
ax.plot(
self.rho,
well,
color=color,
label="DESC M. Landreman eq. 4.19",
)
ax.scatter(self.rho, well, color=color, s=1)
well = self.st["2. DESC Magnetic Well: M. Landreman eq. 3.2 with rho * d/drho"][
1
]
color = MagneticWellVisual._color("2.")
ax.plot(
self.rho,
well,
color=color,
label="DESC M. Landreman eq. 3.2 with rho * d/drho",
)
ax.scatter(self.rho, well, color=color, s=1)
well = self.st["3. DESC Magnetic Well: M. Landreman eq. 3.2"][1]
color = MagneticWellVisual._color("3.")
ax.plot(self.rho, well, color=color, label="DESC M. Landreman eq. 3.2")
ax.scatter(self.rho, well, color=color, s=1)
ax.set(
ylim=[-0.25, 0.25],
xlabel=r"$\rho$",
ylabel="magnetic well",
title=self.name + " well zoomed",
facecolor="white",
)
ax.axhline(color="tab:red")
ax.grid()
fig.legend(fontsize="xx-small")
fig.savefig(self.name + " magnetic wells sym default.png")
torus = MagneticWellVisual("torus", Equilibrium())
dshape = MagneticWellVisual("DSHAPE")
heliotron = MagneticWellVisual("HELIOTRON")
atf = MagneticWellVisual("ATF")
axisym = MagneticWellVisual("AXISYM")
stellarators = (torus, dshape, heliotron, atf, axisym)
/home/kaya/Documents/edu/pton/plasma/DESC/desc/configuration.py:344: UserWarning: Must specify either iota or current. Using default profile of iota=0. warnings.warn( /home/kaya/Documents/edu/pton/plasma/DESC/desc/io/hdf5_io.py:108: RuntimeWarning: Save attribute '_current' was not loaded. warnings.warn(
# to make sure equilibrium were solved correctly on my computer
for s in stellarators:
plot_surfaces(s.eq)
for s in stellarators:
s.print_values()
torus Equilibrium at 0x7f46d93d7700 (L=1, M=1, N=0, NFP=1, sym=False, spectral_indexing=ansi) Precomputing transforms grid.spacing(dr,dt,dz) [1.0000 1.2566 6.2832] 1. DESC Magnetic Well: M. Landreman eq. 4.19 [0.0000] 2. DESC Magnetic Well: M. Landreman eq. 3.2 with rho * d/drho [-3.2159e-16] 3. DESC Magnetic Well: M. Landreman eq. 3.2 [-1.6080e-16] B square average [0.1013] d(magnetic pressure average)/drho [-3.2584e-17] d(thermal pressure)/drho [0.0000] d(total pressure average)/drho [-3.2584e-17] d(total pressure average)/dvolume [-8.2537e-20] d^2(volume)/d(rho)^2 [394.7842] dtdz [39.4784] dvolume/drho [394.7842] volume [197.3921] data["V"] 197.39208802178715 DSHAPE Equilibrium at 0x7f46d92fbd30 (L=26, M=13, N=0, NFP=1.0, sym=1, spectral_indexing=fringe) Precomputing transforms grid.spacing(dr,dt,dz) [1.2599 0.1466 7.9163] 1. DESC Magnetic Well: M. Landreman eq. 4.19 [0.0000] 2. DESC Magnetic Well: M. Landreman eq. 3.2 with rho * d/drho [0.0629] 3. DESC Magnetic Well: M. Landreman eq. 3.2 [0.0369] B square average [0.0575] d(magnetic pressure average)/drho [0.0036] d(thermal pressure)/drho [0.0000] d(total pressure average)/drho [0.0036] d(total pressure average)/dvolume [2.1337e-05] d^2(volume)/d(rho)^2 [305.9274] dtdz [39.4784] dvolume/drho [169.7811] volume [99.4570] data["V"] 99.1767790025823 HELIOTRON Equilibrium at 0x7f46d8fc2760 (L=24, M=12, N=3, NFP=19.0, sym=1, spectral_indexing=fringe) Precomputing transforms grid.spacing(dr,dt,dz) [3.3620 0.5559 0.0855] 1. DESC Magnetic Well: M. Landreman eq. 4.19 [0.0000] 2. DESC Magnetic Well: M. Landreman eq. 3.2 with rho * d/drho [-2.2638] 3. DESC Magnetic Well: M. Landreman eq. 3.2 [-0.9240] B square average [0.1253] d(magnetic pressure average)/drho [-0.2837] d(thermal pressure)/drho [0.0000] d(total pressure average)/drho [-0.2837] d(total pressure average)/dvolume [-0.0006] d^2(volume)/d(rho)^2 [1483.7368] dtdz [39.4784] dvolume/drho [440.0806] volume [179.6268] data["V"] 180.78725668914205 ATF Equilibrium at 0x7f46d8f25b80 (L=24, M=12, N=4, NFP=12.0, sym=1, spectral_indexing=fringe) Precomputing transforms grid.spacing(dr,dt,dz) [2.8845 0.5331 0.1162] 1. DESC Magnetic Well: M. Landreman eq. 4.19 [0.0000] 2. DESC Magnetic Well: M. Landreman eq. 3.2 with rho * d/drho [-1.3759] 3. DESC Magnetic Well: M. Landreman eq. 3.2 [-0.5825] B square average [16.7004] d(magnetic pressure average)/drho [-22.9789] d(thermal pressure)/drho [0.0000] d(total pressure average)/drho [-22.9789] d(total pressure average)/dvolume [-3.1056] d^2(volume)/d(rho)^2 [18.0245] dtdz [39.4784] dvolume/drho [7.3991] volume [3.1325] data["V"] 3.1505967937703065 AXISYM Equilibrium at 0x7f46d93d7d00 (L=40, M=20, N=0, NFP=1.0, sym=1, spectral_indexing=fringe) Precomputing transforms grid.spacing(dr,dt,dz) [1.2599 0.0965 7.9163] 1. DESC Magnetic Well: M. Landreman eq. 4.19 [0.0000] 2. DESC Magnetic Well: M. Landreman eq. 3.2 with rho * d/drho [-0.1511] 3. DESC Magnetic Well: M. Landreman eq. 3.2 [-0.0958] B square average [0.0683] d(magnetic pressure average)/drho [-0.0103] d(thermal pressure)/drho [0.0000] d(total pressure average)/drho [-0.0103] d(total pressure average)/dvolume [-6.6500e-05] d^2(volume)/d(rho)^2 [493.9344] dtdz [39.4784] dvolume/drho [155.3259] volume [98.4592] data["V"] 98.19957214489737
Although the symmetry bug is fixed, forcing it off results in a closer match between the STELLOPT data and DESC 4.19 well. Here symmetry is not forced to false because I think its close enough.
for s in stellarators:
s.compute_plot()
# compare volume computation to known analytic case for torus
fig, ax = plt.subplots()
y = "volume"
rho, volume = torus.st[y]
ax.scatter(rho, volume, label=y + " computation", s=10)
ax.plot(rho, 20 * (np.pi * rho) ** 2, label="analytic")
ax.set(xlabel=r"$\rho$", ylabel=y, title=torus.name)
ax.grid()
fig.legend()
<matplotlib.legend.Legend at 0x7f46d7364be0>
for s in stellarators:
s.plot()
s.plot_magnetic_wells()
Turn symmetry off for closer match.
stellopt_rho, stellopt_well = heliotron.st["0. STELLOPT Magnetic Well"]
desc_well = heliotron.st["1. DESC Magnetic Well: M. Landreman eq. 4.19"][1]
fig, ax = plt.subplots(dpi=600)
ax.plot(stellopt_rho, stellopt_well, label="STELLOPT")
ax.plot(heliotron.rho, desc_well, label="DESC")
ax.scatter(stellopt_rho, stellopt_well, s=1)
ax.scatter(heliotron.rho, desc_well, s=1)
ax.set(
yscale="symlog",
xlabel=r"$\rho$",
ylabel="magnetic well",
title=heliotron.name + " magnetic well",
facecolor="white",
)
ax.axhline(color="tab:red")
ax.grid()
fig.legend()
fig.savefig("Heliotron STELLOPT vs DESC 4.19 sym true.png")
for s in stellarators:
s.save()
In MagneticWell.compute(), the volume computation relies on the quantity d$\theta$ * d$\zeta$. This is obtained from splicing grid.spacing. The grid class takes advantage of things like reducing duplicated node weight, stellarator symmetry, and the number of field periods to reduce the amount of computation.
The problem is that, in the grid class's current implementation, the modifications made to grid.spacing to take advantage of symmetry and NFP do not preserve the value of d$\theta$ d$\zeta$. So independently plucking out d$\theta$ d$\zeta$ from a grid will not give the value its name implies.
Why is this? Because rescaling is done to preserve the full volume, or the weights (d$\rho$ d$\theta$ d$\zeta$) so that grid.spacing.prod(axis=1).sum() = 4$\pi$2, at the expense of changing d$\theta$ * d$\zeta$. Here's a visual:
# For any constant rho surface and NFP=1, grid.spacing is going
# to be an M*N length stack of the row vector (1, 2pi/M, 2pi/N).
g1 = LinearGrid(M=2, N=2, rho=np.array(1))
print(g1.spacing)
print()
# When NFP != 1, the values of drho are increased (decreased) while
# dtheta*dzeta decreases (increases).
# That's why the wrong volume was being computed for the heliotron with NFP > 1.
NFP = np.random.random_sample() * 100
print("NFP:", NFP)
g2 = LinearGrid(M=2, N=2, NFP=NFP, rho=np.array(1))
print(g2.spacing)
# the good thing is that the differential volume drho*dtheta*dzeta is preserved
# since that was the (main?) goal of the rescaling
assert np.allclose(g1.spacing.prod(axis=1), g2.spacing.prod(axis=1))
# and that, because rho=constant forces the drho vector to = 1 (with NFP=1),
# grid.weights is actually the value we expect from dtheta*dzeta (for any NFP).
if g1.num_rho == 1:
assert np.allclose(g1.spacing.prod(axis=1), g1.spacing[:, 1:].prod(axis=1))
assert np.allclose(g1.spacing.prod(axis=1), g1.weights) # definition of weights
[[1.0000 1.2566 1.2566] [1.0000 1.2566 1.2566] [1.0000 1.2566 1.2566] [1.0000 1.2566 1.2566] [1.0000 1.2566 1.2566] [1.0000 1.2566 1.2566] [1.0000 1.2566 1.2566] [1.0000 1.2566 1.2566] [1.0000 1.2566 1.2566] [1.0000 1.2566 1.2566] [1.0000 1.2566 1.2566] [1.0000 1.2566 1.2566] [1.0000 1.2566 1.2566] [1.0000 1.2566 1.2566] [1.0000 1.2566 1.2566] [1.0000 1.2566 1.2566] [1.0000 1.2566 1.2566] [1.0000 1.2566 1.2566] [1.0000 1.2566 1.2566] [1.0000 1.2566 1.2566] [1.0000 1.2566 1.2566] [1.0000 1.2566 1.2566] [1.0000 1.2566 1.2566] [1.0000 1.2566 1.2566] [1.0000 1.2566 1.2566]] NFP: 43.548601624052516 [[3.5182 4.4211 0.1015] [3.5182 4.4211 0.1015] [3.5182 4.4211 0.1015] [3.5182 4.4211 0.1015] [3.5182 4.4211 0.1015] [3.5182 4.4211 0.1015] [3.5182 4.4211 0.1015] [3.5182 4.4211 0.1015] [3.5182 4.4211 0.1015] [3.5182 4.4211 0.1015] [3.5182 4.4211 0.1015] [3.5182 4.4211 0.1015] [3.5182 4.4211 0.1015] [3.5182 4.4211 0.1015] [3.5182 4.4211 0.1015] [3.5182 4.4211 0.1015] [3.5182 4.4211 0.1015] [3.5182 4.4211 0.1015] [3.5182 4.4211 0.1015] [3.5182 4.4211 0.1015] [3.5182 4.4211 0.1015] [3.5182 4.4211 0.1015] [3.5182 4.4211 0.1015] [3.5182 4.4211 0.1015] [3.5182 4.4211 0.1015]]
On grids which are defined completely by a single rho surface (grid.num_rho = 1), the correct volume will be computed $\forall NFP \in \mathbb{R}^{+}$, if we take d$\theta$ * d$\zeta$ to be grid.weights instead of grid.spacing[:, 1:].prod(axis=1).
On grids which are defined over a range of rho surfaces (grid.num_rho != 1) this will not work because d$\rho$ != 1 on those grids. A general solution follows for computing the correct value for d$\theta$ * d$\zeta$ to other grids $g \in Grid : g.num\_rho \in \mathbb{N}$.
g2.weights / g1.spacing[:, 0].grid.spacing) will yield incorrect quantities $\forall g \in Grid : g.num\_rho \in \mathbb{N} \setminus {1}\ and\ g.NFP \in \mathbb{R}^{+} \setminus {1}$.The last source of error in the volume is the stellarator symmetry boolean. enforce_symmetry() again modifies d$\theta$ * d$\zeta$. Visual given below.
The error in computed volume that results from this is much smaller than the NFP error fixed above. The nodes with $\theta$ > $\pi$ are removed and so too are the spacings in the grid that correspond to those nodes. Now because grid.spacing is missing these nodes, to preserve the overall volume in ($\rho$, $\theta$, $\zeta$) space (grid.weights.sum()), the differential volume d$\rho$ d$\theta$ d$\zeta$ of the remaining spaces is increased so that the sum is still 4$\pi$2. This makes sense since we do want to increase the weight of the nodes 0 to $\pi$ to double count for the removed nodes. This process has preserving the volume d$\rho$ d$\theta$ d$\zeta$ in mind. It's not clear to me if a different scale factor (other than 4$\pi$2) would be better to preserve an area element d$\theta$ * d$\zeta$. Might be dependent on $\rho$.
In any case, the solution implemented here is to force stellarator symmetry to False. This is not ideal since it can simplify the computation of the transforms and such. An alternative could be to just recreate the grid with the nodes that were removed from enforce_symmetry() in the MagneticWell class and use this grid just for the volume computation. This could allow the grid with symmetry to still be used elsewhere.
g3 = LinearGrid(M=2, N=2, sym=True, rho=np.array(1))
print(g3.spacing)
[[1.2599 1.3194 1.5833] [1.2599 1.3194 1.5833] [1.2599 1.3194 1.5833] [1.2599 1.3194 1.5833] [1.2599 1.3194 1.5833] [1.2599 1.3194 1.5833] [1.2599 1.3194 1.5833] [1.2599 1.3194 1.5833] [1.2599 1.3194 1.5833] [1.2599 1.3194 1.5833] [1.2599 1.3194 1.5833] [1.2599 1.3194 1.5833] [1.2599 1.3194 1.5833] [1.2599 1.3194 1.5833] [1.2599 1.3194 1.5833]]